﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using NeuralNetworkPro;
using MathNet.Numerics.LinearAlgebra.Double;
using MathNet.Numerics.LinearAlgebra.Double.Factorization;

namespace NeuralNetworkPro
{
    public class Pruning
    {
        private NeuralNetwork net;

        public NeuralNetwork Net
        {
            get { return net; }
            private set { net = value; }
        }

        private double[][] trainTest;

        public double[][] TrainTest
        {
            get { return trainTest; }
            set { trainTest = value; }
        }        

        public Pruning(NeuralNetwork net, double[][] train)
        {
            this.net = net;
            this.trainTest = train;
        }

        public DenseVector[][] ComputeBs()
        {
            DenseVector[][] bs = new DenseVector[net.LayersCount - 2][];
            for (int l = 1; l < net.LayersCount - 1; l++)
            {
                var layer = net[l];
                var front_layer = net[l + 1];

                bs[l - 1] = new DenseVector[layer.NeuronsCount];

                for (int i = 0; i < layer.NeuronsCount; i++)
                {
                    if (layer.NeuronsCount - 1 > 0)
                    {
                        bs[l - 1][i] = new DenseVector(front_layer.NeuronsCount * trainTest.Length);
                    }
                }
            }

            for (int u = 0; u < trainTest.Length; u++)
            {
                net.Compute(trainTest[u]);
                for (int l = 1; l < net.LayersCount - 1; l++)
                {
                    var layer = net[l];
                    var front_layer = net[l + 1];
                    for (int h = 0; h < layer.NeuronsCount; h++)
                    {
                        var b = bs[l - 1][h];
                        if (b == null)
                            continue;
                        for (int i = 0; i < front_layer.NeuronsCount; i++)
                        {
                            var idx = u * front_layer.NeuronsCount + i;
                            var w_ih = front_layer[i][h];
                            var y_hu = layer[h].Output;
                            b[idx] = w_ih * y_hu;
                        }
                    }
                }
            }
            
            return bs;
        }

        private DenseMatrix ComputeA(int layer_id, int h)
        {
            var layer = net[layer_id + 1];
            var front_layer = net[layer_id + 2];
            DenseMatrix a = new DenseMatrix(front_layer.NeuronsCount * trainTest.Length,
                                            front_layer.NeuronsCount * (layer.NeuronsCount - 1));

            for (int u = 0; u < trainTest.Length; u++)
            {
                net.Compute(trainTest[u]);
                for (int i = 0; i < front_layer.NeuronsCount; i++)
                {
                    var idx = u * front_layer.NeuronsCount + i;
                    for (int j = 0; j < layer.NeuronsCount; j++)
                    {
                        if (j != h)
                        {
                            var other_idx =
                                (i * (layer.NeuronsCount - 1)) + (j < h ? j : (j - 1));
                            var y_ju = layer[j].Output;
                            //Console.WriteLine(h + " ping... " + idx + "; " + other_idx);
                            a[idx, other_idx] = y_ju;
                        }
                    }
                }
            }

            return a;
        }

        public void Prune()
        {
            var bs = ComputeBs();
            var indexes = bs.IndexOfMin(t => t.Norm(2));
            var layer_id = indexes.Item1;
            var index_h = indexes.Item2;
            var a = ComputeA(layer_id, index_h);
            var b = bs[layer_id][index_h];
            var deltas = SolveLeastSquares(a, b);

            net.RemoveNeuron(layer_id + 1, index_h);

            UpdateWeights(layer_id + 1, deltas);

            Console.WriteLine("index_h");


        }

        private void UpdateWeights(int layer_id, DenseVector deltas)
        {
            var layer = net[layer_id];
            var front_layer = net[layer_id + 1]; 
            for (int i = 0; i < front_layer.NeuronsCount; i++)
            {
                for (int j = 0; j < layer.NeuronsCount; j++)
                {
                    var other_idx =
                            (i * layer.NeuronsCount) + j;
                    var d = deltas[other_idx];
                    front_layer[i][j] += d;

                }
            }
        }



        private DenseVector SolveLeastSquares(DenseMatrix a, DenseVector b)
        {
            var y = a;
            var z = b;
            Console.WriteLine("Applying CGPCNE");
            var deltas = ComputeCGPCNE(y, z, 1.0, 0.01);

            #region old

            ////////////////var m = a.RowCount;
            ////////////////var n = a.ColumnCount;

            ////////////////var beta = (a.TransposeThisAndMultiply(a));
            //////////////////Console.WriteLine(beta);
            ////////////////Console.WriteLine("rank = " + beta.Rank());
            ////////////////beta = beta.Inverse().TransposeAndMultiply(a);
            //////////////////Console.WriteLine(beta);
            ////////////////var res = beta.Multiply(b);
            ////////////////var x = res.Column(0);
            ////////////////if (double.IsNaN(x.Sum()))
            ////////////////    throw new Exception("Matrix is singular");
            ////////////////Console.WriteLine(res[0, 0]);

            //////////////////// compute the SVD
            //////////////////Svd svd = a.Svd(true);

            //////////////////// get matrix of left singular vectors with first n columns of U
            //////////////////var U1 = svd.U().SubMatrix(0, m, 0, n);
            //////////////////// get matrix of singular values
            //////////////////var S = new DiagonalMatrix(n, n, svd.S().ToArray());
            //////////////////// get matrix of right singular vectors
            //////////////////var V = svd.VT().Transpose();

            //////////////////var x = V.Multiply(S.Inverse()).Multiply(U1.Transpose().Multiply(b));
            //////////////////Console.WriteLine(x);

            #endregion

            return deltas;
        }

        /// <summary>
        /// Conjugate gradient preconditioned normal equation
        /// </summary>
        /// <param name="y"></param>
        /// <param name="z"></param>
        /// <param name="w"></param>
        /// <param name="epsilon"></param>
        /// <returns></returns>
        private DenseVector ComputeCGPCNE(DenseMatrix y, DenseVector z, double w, double epsilon)
        {
            var k = y.ColumnCount;
            var d = ComputeD(y);
            //var l = ComputeL(y);
            //var c_w = ComputeCW(d, l, w);
            //var c_w_i = c_w.Inverse();
            //var c_w_i_t = c_w_i.Transpose();
            //var y_t = y.Transpose();


            var delta_old = new DenseVector(y.ColumnCount);
            var delta_new = new DenseVector(y.ColumnCount);
            var r = DenseVector.OfVector(z - (y.Multiply(delta_old)));
            //var s = c_w_i.Multiply(y_t) * r; 
            var s = ComputeS(y, d, DenseVector.OfVector(r), w);
            var p = DenseVector.OfVector(s);
            do
            {
                {
                    var delta_tmp = delta_old;
                    delta_old = delta_new;
                    delta_new = delta_tmp;
                }
                var t_q = ComputeTQ(y, d, p, w);
                var t = t_q.Item1; //c_w_i_t.Multiply(p);
                var q = t_q.Item2; //y.Multiply(t); 

                var a = Math.Pow(s.Norm(2), 2) / Math.Pow(q.Norm(2), 2);
                r = r - a * q;
                var s_new = ComputeS(y, d, r, w); //c_w_i.Multiply(y_t).Multiply(r);
                var b = Math.Pow(s_new.Norm(2), 2) / Math.Pow(s.Norm(2), 2);
                var p_new = s_new + b * p;
                delta_new = DenseVector.OfVector(delta_old + a * t);
                s = s_new;
                p = DenseVector.OfVector(p_new);
                //Console.WriteLine("At step: " + step++ + " error = " + (delta_old - delta_new).Norm(2) + " compared with " + epsilon + " that is " + ((delta_old - delta_new).Norm(2) < epsilon));
                //ites++;
            }
            while ((delta_old - delta_new).Norm(2) >= epsilon);
            
            return delta_new;
        }

        private static DiagonalMatrix ComputeD(DenseMatrix y)
        {
            var k = y.ColumnCount;
            var d = new DiagonalMatrix(k);
            for (int j = 0; j < k; j++)
            {
                var d_jj = Math.Pow(y.Column(j).Norm(2), 2);
                d[j, j] = d_jj;
            }
            return d;
        }

        private DenseMatrix ComputeL(DenseMatrix y)
        {
            var dim = y.ColumnCount;
            var res = new DenseMatrix(dim);
            for (int k = 0; k < dim; k++)
            {
                for (int j = k + 1; j < dim; j++)
                {
                    var y_j = y.Column(j);
                    var y_k = y.Column(k);
                    res[j, k] = y_j.ScalarProduct(y_k);
                }
            }
            return res;
        }

        private DenseMatrix ComputeCW(DiagonalMatrix d, DenseMatrix l, double w)
        {
            var sum = d + (w * l);
            var pow = d;
            for (int i = 0; i < d.ColumnCount; i++)
            {
                pow[i, i] = Math.Pow(d[i, i], -0.5);
            }
            return DenseMatrix.OfMatrix(sum * pow);
        }



        private Tuple<DenseVector, DenseVector> ComputeTQ(DenseMatrix y, DiagonalMatrix d, DenseVector p, double w)
        {
            // U deve essere altra dim
            var k = y.ColumnCount;

            var t = new DenseVector(k);
            var u = new DenseVector(y.RowCount);

            for (int j = k - 1; j >= 0; j--)
            {
                var scalar = y.Column(j).DotProduct(u); //.ScalarProduct(u);
                t[j] = (Math.Pow(d[j, j], -0.5) * p[j]) - (w * (1 / d[j, j]) * scalar);
                u = u + (t[j] * DenseVector.OfVector(y.Column(j)));
            }

            return new Tuple<DenseVector, DenseVector>(t, u);
        }

        private DenseVector ComputeS(DenseMatrix y, DiagonalMatrix d, DenseVector r, double w)
        {
            var k = y.ColumnCount;
            var s = new DenseVector(k);
            var u = r;
            for (int j = 0; j < k; j++)
            {
                var y_j = y.Column(j);
                var d_j_pow = Math.Pow(d[j, j], -0.5);
                s[j] = d_j_pow * (y.Column(j).DotProduct(u)); //.ScalarProduct(u));
                u = u - (w * d_j_pow * s[j]) * DenseVector.OfVector(y_j);
            }
            return s;
        }

        public bool CanPrune()
        {
            return net.HiddenLayers.Any(l => l.NeuronsCount > 1);
        }
    }
}
